import librosa
import numpy as np
from sklearn.svm import SVC
import os

def extract_mfcc_from_folder(folder_path):
    """
    从给定文件夹中提取所有音频文件的MFCC特征
    """
    all_mfcc = []
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            if file.lower().endswith(('.wav', '.mp3')):  # 确保是音频文件
                file_path = os.path.join(root, file)
                mfcc = extract_mfcc(file_path)
                all_mfcc.append(mfcc)
    return all_mfcc

def extract_mfcc(file_path):
    y, sr = librosa.load(file_path, sr=None)
    mfcc = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13)
    mfcc = np.mean(mfcc.T, axis=0)
    return mfcc

# 训练集文件夹路径，里面分别有'你'和'好'的子文件夹，子文件夹中存放对应音频
train_folder_path = r'E:\实训\音频数据集'
# 预测的音频所在文件夹路径
predict_folder_path = r'E:\实训\测试数据集'

# 提取训练集的MFCC特征
mfcc_ni = extract_mfcc_from_folder(os.path.join(train_folder_path, '你'))
mfcc_hao = extract_mfcc_from_folder(os.path.join(train_folder_path, '好'))

# 特征向量和标签
X = mfcc_ni + mfcc_hao  # 特征向量
y = ['你'] * len(mfcc_ni) + ['好'] * len(mfcc_hao)  # 标签

# 构建 SVM 模型
svm = SVC(kernel='linear')
svm.fit(X, y)

# 提取预测文件夹中音频的MFCC特征并进行预测
for root, dirs, files in os.walk(predict_folder_path):
    for file in files:
        if file.lower().endswith(('.wav', '.mp3')):  # 确保是音频文件
            file_path = os.path.join(root, file)
            mfcc_to_predict = extract_mfcc(file_path)
            prediction = svm.predict([mfcc_to_predict])
            print(f'Prediction for {file}: {prediction[0]}')